package flink.datastream;

import flink.model.Order;
import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.tuple.Tuple1;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.junit.Test;

import java.util.Arrays;

/**
 * @author binarylei
 * @version 2019-11-28
 */
public class PartitionDataStreamTest {

    @Test
    public void testShuffle() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        DataStreamSource<Order> ds1 = env.fromCollection(Arrays.asList(
                new Order(1L, "order1"),
                new Order(2L, "order2")));

        ds1.shuffle().print().setParallelism(3);
        env.execute();
    }

    @Test
    public void testRebalance() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        DataStreamSource<Order> ds1 = env.fromCollection(Arrays.asList(
                new Order(1L, "order1"),
                new Order(2L, "order2")));

        ds1.rebalance().print().setParallelism(3);
        env.execute();
    }

    @Test
    public void testMyPartition() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        DataStreamSource<Tuple1> ds1 = env.fromCollection(Arrays.asList(
                new Tuple1<Long>(1L),
                new Tuple1<Long>(2L),
                new Tuple1<Long>(3L)));

        ds1.partitionCustom(new MyPartition(), 0).print().setParallelism(3);
        env.execute();
    }

    // 流广播分区，每个分区都会广播一份数据
    @Test
    public void testBroadcast() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        DataStreamSource<Tuple1> ds1 = env.fromCollection(Arrays.asList(
                new Tuple1<Long>(1L)));

        // 3个并行度，则输出3次日志
        ds1.broadcast().print().setParallelism(3);
        env.execute();
    }

    public class MyPartition implements Partitioner<Long> {

        @Override
        public int partition(Long aLong, int i) {
            return (int) (aLong % 2);
        }
    }
}
